// #include <stdio.h>
// #include <NvInfer.h>
// #include <NvOnnxParser.h>
// #include <cuda.h>
// #include <cuda_runtime.h>
// #include <iostream>
// #include <fstream>

// using namespace nvinfer1;
// using namespace nvonnxparser;

// class Logger : public ILogger
// {
//     void  log(Severity severity, const char* msg) noexcept
//     {
//         if (severity <= Severity::kERROR)
//         {
//             std::cout << msg << std::endl;
//         }
//     }
// };

// int main()
// {
//     Logger logger;
//     printf("hello, world\n");
//     IBuilder *m_builder =  createInferBuilder(logger);
//     m_builder->getLogger()->log(nvinfer1::ILogger::Severity::kERROR, "Create builder...");

//     uint32_t flag = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);

//     IBuilderConfig *m_buildConfig = m_builder->createBuilderConfig();
//     INetworkDefinition *m_network = m_builder->createNetworkV2(flag);
//     // IParser* m_parser = createNvOnnxParser_INTERNAL(m_network, &logger);
//     IParser *m_parser = createParser(*m_network, logger);
//      if(!m_parser->parseFromFile("/home/xtdz/project/model/yolov5s.onnx", static_cast<int>(ILogger::Severity::kINTERNAL_ERROR)))
//      {
//         std::cout << "Error: unable to parase onnx model" << std::endl;
//      }

//       ICudaEngine *m_cudaEngine = m_builder->buildEngineWithConfig(*m_network, *m_buildConfig);
//       IHostMemory* serializeModel = m_cudaEngine->serialize();

//       std::ofstream engineFileOut("/home/xtdz/project/model/yolov5s.engine", std::ios::binary);

//       engineFileOut.write(static_cast<const char*>(serializeModel->data()), serializeModel->size());

//       engineFileOut.close();

//     delete m_parser;
//     delete m_network;
//     delete m_buildConfig;
//     delete m_builder;
//     delete serializeModel;
    
//     return 0;
// }

#include <iostream>
#include <fstream>
#include <string>
#include <NvInfer.h>
#include <NvOnnxParser.h>

using namespace nvinfer1;
using namespace nvonnxparser;

bool fileExists(const std::string& filePath) {
    std::ifstream file(filePath);
    return file.good();
}

class Logger : public ILogger
{
    void  log(Severity severity, const char* msg) noexcept
    {
        if (severity <= Severity::kERROR)
        {
            std::cout << msg << std::endl;
        }
    }
};

int main(int argc, char** argv) {

    if (argc < 3)
    {
        std::cout << "usage: ./exec input_onnx outputEngine" << std::endl;
        return 0;
    }
    

    // std::string onnxModelPath = "/home/xtdz/project/model/yolov5s.onnx";
    // std::string engineFilePath = "/home/xtdz/project/model/yolov5s.engine";
    std::string onnxModelPath = argv[1];
    std::string engineFilePath = argv[2];

    if (!fileExists(onnxModelPath)) {
        std::cerr << "ONNX model file not found: " << onnxModelPath << std::endl;
        return -1;
    }

    // Initialize logger
    Logger logger;
    IBuilder* builder = createInferBuilder(logger);
    uint32_t flag = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);

    // Create network and parser
    INetworkDefinition* network = builder->createNetworkV2(flag);
    // INetworkDefinition* network = builder->createNetworkV2(0U);
    IParser* parser = createParser(*network, logger);

    if (!parser->parseFromFile(onnxModelPath.c_str(), static_cast<int>(ILogger::Severity::kERROR))) {
        std::cerr << "Error: Unable to parse ONNX model: " << onnxModelPath << std::endl;
        return -1;
    }

    IBuilderConfig* config = builder->createBuilderConfig();

    // config->setFlag(BuilderFlag::kFP16);
    // config->setFlag(BuilderFlag::kGPU_FALLBACK);
    // config->setDefaultDeviceType(DeviceType::kDLA);
    // config->setDLACore(0);

    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    IHostMemory* serializedModel = engine->serialize();

    // Save engine to file
    std::ofstream engineFileOut(engineFilePath, std::ios::binary);
    engineFileOut.write(static_cast<const char*>(serializedModel->data()), serializedModel->size());
    engineFileOut.close();

    // Cleanup
    serializedModel->destroy();
    engine->destroy();
    config->destroy();
    network->destroy();
    parser->destroy();
    builder->destroy();

    std::cout << "Engine file created successfully: " << engineFilePath << std::endl;
    return 0;
}

